# local_drone_publisher.py (运行在你的本机 Ubuntu)

import paho.mqtt.client as mqtt
import json
import time
import uuid # 用于生成唯一的请求ID

# --- MQTT 配置 ---
MQTT_BROKER_HOST = "118.89.105.38"
MQTT_BROKER_PORT = 1883
MQTT_USERNAME = "buaauav"
MQTT_PASSWORD = "buaauav"

# MQTT 话题定义
# 这个话题将模拟 /mavros/setpoint_raw/local
SIMULATED_ROS_TOPIC = "drone/trajectory/setpoint_raw"
COMMAND_TOPIC = "drone/command/trigger_algo" # 用于接收远端指令

CLIENT_ID = "local_drone_publisher_001"

# --- 全局变量 ---
# 用于控制是否开始发布模拟航迹数据
is_publishing_trajectory = False
current_request_id = None # 用于关联指令和发布的数据

# --- MQTT 回调函数 ---
def on_connect(client, userdata, flags, rc):
    if rc == 0:
        print(f"MQTT Connected to Broker ({MQTT_BROKER_HOST})!")
        client.subscribe(COMMAND_TOPIC) # 订阅命令话题
        print(f"Subscribed to command topic: {COMMAND_TOPIC}")
    else:
        print(f"MQTT Failed to connect, return code {rc}")
        exit(1)

def on_message(client, userdata, msg):
    global is_publishing_trajectory, current_request_id
    try:
        payload = json.loads(msg.payload.decode())
        command = payload.get("command")
        request_id = payload.get("request_id")

        if command == "generate_trajectory":
            print(f"Received command: generate_trajectory for Request ID: {request_id}")
            is_publishing_trajectory = True
            current_request_id = request_id
        elif command == "stop_trajectory":
            print(f"Received command: stop_trajectory for Request ID: {request_id}")
            is_publishing_trajectory = False
            current_request_id = None
        else:
            print(f"Received unknown command: {command}")
    except json.JSONDecodeError:
        print(f"Received non-JSON message on {msg.topic}: {msg.payload.decode()}")
    except Exception as e:
        print(f"Error processing MQTT message: {e}")

# --- 主程序逻辑 ---
if __name__ == "__main__":
    # 初始化 MQTT 客户端
    client = mqtt.Client(client_id=CLIENT_ID)
    if MQTT_USERNAME and MQTT_PASSWORD:
        client.username_pw_set(MQTT_USERNAME, MQTT_PASSWORD) # 设置用户名和密码

    client.on_connect = on_connect
    client.on_message = on_message

    try:
        client.connect(MQTT_BROKER_HOST, MQTT_BROKER_PORT, 60)
        client.loop_start() # 在后台线程运行，处理消息收发
    except Exception as e:
        print(f"Failed to connect to MQTT broker: {e}")
        exit(1)

    print("Simulated Drone Publisher started. Waiting for commands...")

    # 模拟航迹生成和发布
    i = 0
    while True:
        if is_publishing_trajectory:
            i += 1
            # 模拟 PositionTarget ROS 消息的 JSON 数据结构
            # 这里的字段名和结构尽可能与 mavros_msgs/PositionTarget 保持一致，以便未来切换
            simulated_setpoint_raw_data = {
                "request_id": current_request_id,
                "header": {
                    "stamp": time.time(), # 使用当前时间戳
                    "frame_id": "map"
                },
                "coordinate_frame": 1, # 模拟 PositionTarget.FRAME_LOCAL_NED
                "type_mask": 0, # 模拟 type_mask
                "x": 10.0 + i * 0.1,
                "y": 5.0 + i * 0.05,
                "z": 2.0,
                "vx": 0.1,
                "vy": 0.05,
                "vz": 0.0,
                "afx": 0.0, "afy": 0.0, "afz": 0.0,
                "yaw": 0.0, "yaw_rate": 0.0
            }
            # 将字典转换为 JSON 字符串
            payload_json = json.dumps(simulated_setpoint_raw_data)

            try:
                # 发布到 MQTT 话题
                client.publish(SIMULATED_ROS_TOPIC, payload_json, qos=1)
                print(f"Published point {i} for Request ID {current_request_id} to MQTT topic: {SIMULATED_ROS_TOPIC}")
            except Exception as e:
                print(f"Failed to publish MQTT message: {e}")

        time.sleep(0.5) # 每0.5秒发布一个模拟点

    client.loop_stop()
    client.disconnect()
    print("Publisher stopped.")
